package LibDL.eval.classification;

import LibDL.core.Dtype;
import LibDL.core.Scalar;
import LibDL.core.Tensor;
import LibDL.core.functional;

public class FScoreEvaluator<T> extends AbstractClassificationEvaluator<T> {
    private double beta;
    private PRFReduction reduction;
    @Override
    protected double core(Tensor truth, Tensor pred) {
        if (getBeta() <= 0){
            throw new IllegalArgumentException("beta should be >0 for F-beta score, Found: "+getBeta());
        }
        Tensor mcm = setSampleWise(getReduction().equals(PRFReduction.SAMPLES))
                .getMultilabelConfusionMatrix();
        if (mcm.size(0) != 1 && getReduction().equals(PRFReduction.BINARY)){
            throw new IllegalArgumentException("\"binary\" is not supported for current target.");
        }

        Tensor tpSum = mcm.select(1,0).select(1,0);
        Tensor trueSum = mcm.select(1,0).sum(1);
        Tensor predSum = mcm.select(2,0).sum(1);

        if (getReduction().equals(PRFReduction.MICRO)){
            tpSum = tpSum.sum().unsqueeze(0);
            trueSum = trueSum.sum().unsqueeze(0);
            predSum = predSum.sum().unsqueeze(0);
        }
        // avoid any/0
        Tensor precisionZeroDivPoints = functional.equal_indices(predSum,new Scalar(0.0)).to(Dtype.FLOAT64);
        Tensor recallZeroDivPoints = functional.equal_indices(trueSum,new Scalar(0.0)).to(Dtype.FLOAT64);
        predSum = predSum.add(precisionZeroDivPoints);
        trueSum = trueSum.add(recallZeroDivPoints);
        Tensor precision = tpSum.div(predSum);
        Tensor recall = tpSum.div(trueSum);

        Tensor denom = precision.mul(new Scalar(getBeta()*getBeta())).add(recall);
        // avoid /0
        Tensor fScoreZeroDivPoints = functional.equal_indices(denom, new Scalar(0.0)).to(Dtype.FLOAT64);
        denom = denom.add(fScoreZeroDivPoints);
        Tensor fScore = precision.mul(recall).mul(new Scalar(getBeta()*getBeta()+1.0)).div(denom);

        fScore = fScore.add(fScoreZeroDivPoints.mul(new Scalar(getZeroDivValue())));
        Tensor weight;
        if (getReduction().equals(PRFReduction.WEIGHTED)){
            weight = trueSum;
        }else if (getReduction().equals(PRFReduction.SAMPLES)){
            weight = getSampleWeight();
        }else {
            weight = functional.ones(trueSum.size(0));
        }

        return weightedSum(fScore,weight,true,-1).item().to_double();
    }

    public double getBeta() {
        return beta;
    }

    public FScoreEvaluator<T> setBeta(double beta) {
        this.beta = beta;
        return this;
    }

    public PRFReduction getReduction() {
        return reduction;
    }

    public FScoreEvaluator<T> setReduction(PRFReduction reduction) {
        this.reduction = reduction;
        return this;
    }
}
